from ultralytics import YOLO
import cv2
import os
from pathlib import Path

# 加载预训练姿势检测模型
model = YOLO('yolo11n-pose.pt')  # 使用官方预训练模型

# 定义路径
image_dir = os.path.join('datasets', 'images')
label_dir = os.path.join('datasets', 'labels')
os.makedirs(label_dir, exist_ok=True)  # 创建输出目录

# 获取所有图片文件
image_paths = [p for p in Path(image_dir).iterdir() if p.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp']]

for image_path in image_paths:
    # 进行预测
    results = model(image_path)

    # 获取第一个结果（单张图片预测）
    result = results[0]

    # 准备输出文件路径
    txt_path = os.path.join(label_dir, f"{image_path.stem}.txt")

    with open(txt_path, 'w') as f:
        # 遍历每个检测到的人体实例
        for i in range(len(result.boxes)):
            # 获取归一化坐标
            box = result.boxes.xywhn[i].cpu().numpy()  # 框坐标
            cls_id = int(result.boxes.cls[i].item())  # 类别ID
            kpts = result.keypoints.xyn[i].cpu().numpy().flatten()  # 关键点坐标

            # 格式转换
            box_coords = [f"{coord:.6f}" for coord in box]
            kpts_coords = [f"{coord:.6f}" for coord in kpts]

            # 构建输出行
            line = [str(cls_id)] + box_coords + kpts_coords
            f.write(" ".join(line) + "\n")

print("关键点转换完成！")